import json
import re
import spacy
import networkx as nx
from nltk.tokenize import word_tokenize
nlp = spacy.load("en_core_web_sm")

def convert(path_src, path_des):
    with open(path_src, 'r', encoding='utf-8') as fr:  # 读数据
        data = fr.readlines()  # 获取数据集中的每一行
    with open(path_des, 'w', encoding='utf-8') as fw:  # 写数据
        for i in range(0, len(data), 4):
            id_s, sentence = data[i].strip().split('\t')
            sentence = sentence[1:-1]
            sentence = search_entity_dependency(sentence)
            meta = dict(
                id=id_s,
                relation=data[i+1].strip(),
                sentence=sentence,
                comment=data[i+2].strip()[8:]
            )
            json.dump(meta, fw, ensure_ascii=False)
            fw.write('\n')

def search_entity_dependency(sentence):
    # 获取entity1与entity2
    e1 = re.findall(r'<e1>(.*)</e1>', sentence)[0]  # 返回sentence中<e1></e1>之间的内容，返回形式是数组，[0]即实体名e1
    e2 = re.findall(r'<e2>(.*)</e2>', sentence)[0]


    sentence = sentence.replace('<e1>'+e1+'</e1>', 'entity1')
    sentence = sentence.replace('<e2>'+e2+'</e2>', 'entity2')

    doc = nlp(sentence)
    # print('sentence:', format(doc))
    # Load spacy's dependency tree into a networkx graph
    edges = []
    for token in doc:
        for child in token.children:
            edges.append(('{0}'.format(token),
                          '{0}'.format(child)))
    graph = nx.Graph(edges)
    # Get the length and path
    # print('shortest path lenth: ', nx.shortest_path_length(graph, source=e1, target=e2))
    # print('shortest path: ', nx.shortest_path(graph, source=e1, target=e2))
    sentence = nx.shortest_path(graph, source='entity1', target='entity2')

    # sentence.insert(0, '<e1>')
    # sentence.insert(2, '</e1>')
    # sentence.insert(len(sentence), '</e2>')
    # sentence.insert(len(sentence) - 2, '<e2>')

    # assert '<e1>' in sentence
    # assert '<e2>' in sentence
    # assert '</e1>' in sentence
    # assert '</e2>' in sentence
    assert 'entity1' in sentence
    assert 'entity2' in sentence

    return sentence

if __name__ == '__main__':
    path_train = './SemEval2010_task8_all_data/SemEval2010_task8_training/TRAIN_FILE.TXT'
    path_test = './SemEval2010_task8_all_data/SemEval2010_task8_testing_keys/TEST_FILE_FULL.TXT'

    # convert(path_train, 'train_dependency.json')
    convert(path_test, 'test_dependency.json')